package net.xt.myupload.service;

import lombok.extern.slf4j.Slf4j;
import net.xt.myupload.util.HttpUtils;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URLConnection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;

@Service
@Slf4j
public class FileDownloader {

    private final int CHUNK_SIZE = 1024 * 1024 * 50;    // 自定义每个文件分片大小 50M
    private final String SAVE_DIR = "D:/result";    // 自定义下载文件保存的路径
    @Resource
    private ThreadPoolTaskExecutor threadPool;     // 创建自定义线程池

    public void downloadFile(String fileURL, String downloadFileName) throws Exception {
        // 第一次请求，只获取文件总大小
        HttpURLConnection conn = HttpUtils.sendPost(fileURL, "fileName=" + downloadFileName);

        int responseCode = conn.getResponseCode();
        if (responseCode == HttpServletResponse.SC_OK) {
            // 获取文件总大小
            long fileTotalSize = conn.getContentLengthLong();
            // 计算总共多少个分片：总大小 / 每个分片大小
            int numChunks = (int) Math.ceil((double) fileTotalSize / CHUNK_SIZE);
            // 定义线程计数器，确保所有线程都执行完成，再进行后续操作
            CountDownLatch countDownLatch = new CountDownLatch(numChunks);

            for (int i = 0; i < numChunks; i++) {
                long startByte = i * CHUNK_SIZE;
                long endByte = Math.min((i + 1) * CHUNK_SIZE - 1, fileTotalSize - 1);
                String chunkFileName = "chunk-" + i + downloadFileName;

                // 使用多线程下载，每个文件分片使用一个线程，提高下载速度和性能
                CompletableFuture.runAsync(() -> {
                    try {
                        // 第二次请求，下载每个分片
                        this.downloadChunk(fileURL, startByte, endByte, downloadFileName, chunkFileName);
                        // 每个线程执行完后，计数器-1
                        countDownLatch.countDown();
                        log.info(chunkFileName + "下载完毕");

                        // 合并所有分片
//                        if (endByte - fileTotalSize >= -1) {
//                            this.mergeChunks(numChunks, downloadFileName);
//                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                        throw new RuntimeException(e);
                    }
                }, threadPool);
            }

            // 等待所有线程执行的分片下载全部完成，再进行合并分片
            countDownLatch.await();
            log.info("所有分片下载完毕，开始合并分片。。。。。。");
            // 合并所有分片
            this.mergeChunks(numChunks, downloadFileName);
            log.info("合并分片完成！！！");
        }
    }

    /**
     * 下载文件分片
     *
     * @param fileURL          下载的url
     * @param startByte        开始字节
     * @param endByte          结束字节
     * @param downloadFileName 下载的文件名
     * @param chunkFileName    文件分片名
     * @throws Exception
     */
    private void downloadChunk(String fileURL,
                               long startByte,
                               long endByte,
                               String downloadFileName,
                               String chunkFileName)
            throws Exception {
        // 构造请求头，每个分片的大小，格式 Range:bytes=startByte-endByte
        Map<String, Object> rangeMap = new HashMap<>();
        Map<String, Long> chunkMap = new HashMap<>();
        chunkMap.put("startByte", startByte);
        chunkMap.put("endByte", endByte);
        rangeMap.put("Range", chunkMap);
        URLConnection conn = HttpUtils.sendPost(fileURL, "fileName=" + downloadFileName, rangeMap);

        try (BufferedInputStream inputStream = new BufferedInputStream(conn.getInputStream());
             FileOutputStream outputStream = new FileOutputStream(SAVE_DIR + "/" + chunkFileName);) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            outputStream.flush();
        }
    }

    /**
     * 合并文件分片
     *
     * @param numChunks        总共的分片数量
     * @param downloadFileName 合并的文件名
     * @throws Exception
     */
    private void mergeChunks(int numChunks, String downloadFileName) throws Exception {
        File resultFile = new File(SAVE_DIR + "/" + downloadFileName);
        try (FileOutputStream outputStream = new FileOutputStream(resultFile)) {
            for (int i = 0; i < numChunks; i++) {
                File fileChunk = new File(SAVE_DIR + "/chunk-" + i + downloadFileName);

                // 由于使用多线程下载，如果分片还没有下载或分片大小不完整，就先短时间等待
                // 前面已经使用CountDownLatch，确保所有线程中的文件分片都下载完成，所以这里可以不要再判断
//                while (!fileChunk.exists() || (i != numChunks - 1 && fileChunk.length() < CHUNK_SIZE)) {
//                    Thread.sleep(100);
//                }

                try (InputStream inputStream = new FileInputStream(fileChunk)) {
                    byte[] buffer = new byte[1024];
                    int bytesRead;
                    while ((bytesRead = inputStream.read(buffer)) != -1) {
                        outputStream.write(buffer, 0, bytesRead);
                    }
                    outputStream.flush();
                }

                // 删除分片
                fileChunk.delete();
            }
            outputStream.flush();
        }
    }

}
